#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project ：V2 
@File    ：DB1.py
@IDE     ：PyCharm 
@Author  ：郭星
@Date    ：2025/9/23 19:14 
'''
import random
import numpy as np
from typing import List, Tuple, Dict


class GasSchedulingGA:
    def __init__(self,
                 num_workpieces: int,
                 num_components: int,
                 num_inflation_eq: int,
                 num_analysis_eq_per_component: int,
                 population_size: int = 50,
                 generations: int = 100,
                 crossover_rate: float = 0.8,
                 mutation_rate: float = 0.1):
        """初始化气体调度遗传算法"""
        self.num_workpieces = num_workpieces
        self.num_components = num_components
        self.num_inflation_eq = num_inflation_eq
        self.num_analysis_eq_per_component = num_analysis_eq_per_component

        self.population_size = population_size
        self.generations = generations
        self.crossover_rate = crossover_rate
        self.mutation_rate = mutation_rate

        # 随机生成加工时间（实际应用中替换为真实数据）
        self.inflation_time = np.random.randint(5, 15, size=num_workpieces)
        self.analysis_time = np.random.randint(3, 10, size=(num_workpieces, num_components))

        self.population = []

    def initialize_chromosome(self) -> Dict:
        """初始化一条染色体"""
        # 充气阶段基因
        inflation_order = random.sample(range(self.num_workpieces), self.num_workpieces)
        inflation_eq = [random.randint(0, self.num_inflation_eq - 1) for _ in range(self.num_workpieces)]

        # 分析阶段基因
        analysis_order = random.sample(range(self.num_workpieces), self.num_workpieces)

        # 组分串联组合基因
        component_groups = []
        group_test_order = []
        analysis_eq = []  # 与group_test_order和组分严格同步
        for _ in range(self.num_workpieces):
            max_groups = random.randint(1, self.num_components)
            groups = []
            for i in range(self.num_components):
                if i < max_groups:
                    groups.append(i)
                else:
                    groups.append(random.randint(0, max_groups - 1))
            component_groups.append(groups)

            # 串联组测试顺序
            unique_groups = list(sorted(set(groups)))
            random.shuffle(unique_groups)
            group_test_order.append(unique_groups)

            # 生成分析设备选择（确保与组内组分数一致）
            eq_for_workpiece = []
            for group_id in unique_groups:
                # 精确计算该组包含的组分数
                components_in_group = [j for j, g in enumerate(groups) if g == group_id]
                # 设备列表长度必须与组分数完全一致
                eq_for_group = [random.randint(0, self.num_analysis_eq_per_component - 1)
                                for _ in components_in_group]
                eq_for_workpiece.append(eq_for_group)
            analysis_eq.append(eq_for_workpiece)

        return {
            'inflation_order': inflation_order,
            'inflation_eq': inflation_eq,
            'analysis_order': analysis_order,
            'component_groups': component_groups,
            'group_test_order': group_test_order,
            'analysis_eq': analysis_eq
        }

    def initialize_population(self):
        """初始化种群"""
        self.population = [self.initialize_chromosome() for _ in range(self.population_size)]

    def calculate_fitness(self, chromosome: Dict) -> float:
        """计算适应度，增加索引有效性检查"""
        # 计算充气阶段时间
        inflation_machine_time = [0] * self.num_inflation_eq
        for workpiece in chromosome['inflation_order']:
            eq = chromosome['inflation_eq'][workpiece]
            start_time = inflation_machine_time[eq]
            end_time = start_time + self.inflation_time[workpiece]
            inflation_machine_time[eq] = end_time
        inflation_completion_time = max(inflation_machine_time)

        # 计算分析阶段时间
        analysis_machine_time = [[0 for _ in range(self.num_analysis_eq_per_component)]
                                 for _ in range(self.num_components)]
        workpiece_inflation_end = {}
        temp_machine_time = [0] * self.num_inflation_eq
        for workpiece in chromosome['inflation_order']:
            eq = chromosome['inflation_eq'][workpiece]
            start_time = temp_machine_time[eq]
            end_time = start_time + self.inflation_time[workpiece]
            temp_machine_time[eq] = end_time
            workpiece_inflation_end[workpiece] = end_time

        # 按分析顺序处理工件
        for workpiece in chromosome['analysis_order']:
            # 检查工件索引有效性
            if workpiece < 0 or workpiece >= self.num_workpieces:
                continue

            start_analysis_time = workpiece_inflation_end.get(workpiece, 0)
            group_test_order = chromosome['group_test_order'][workpiece]
            analysis_eq = chromosome['analysis_eq'][workpiece]

            # 遍历组测试顺序
            for group_idx, group_id in enumerate(group_test_order):
                # 检查组索引有效性
                if group_idx < 0 or group_idx >= len(analysis_eq):
                    continue

                # 获取该组实际包含的组分
                components_in_group = [j for j, g in enumerate(chromosome['component_groups'][workpiece])
                                       if g == group_id]
                # 获取该组的设备列表
                eq_for_group = analysis_eq[group_idx]

                # 关键检查：确保设备列表长度与组分数一致
                if len(eq_for_group) != len(components_in_group):
                    # 修复不一致问题（实际应用中可记录日志）
                    eq_for_group = [random.randint(0, self.num_analysis_eq_per_component - 1)
                                    for _ in components_in_group]
                    analysis_eq[group_idx] = eq_for_group

                # 计算组准备时间
                group_ready_time = start_analysis_time
                for comp_idx, component in enumerate(components_in_group):
                    # 检查组分索引有效性
                    if comp_idx < 0 or comp_idx >= len(eq_for_group):
                        continue
                    eq = eq_for_group[comp_idx]
                    group_ready_time = max(group_ready_time, analysis_machine_time[component][eq])

                # 更新设备时间
                for comp_idx, component in enumerate(components_in_group):
                    if comp_idx < 0 or comp_idx >= len(eq_for_group):
                        continue
                    eq = eq_for_group[comp_idx]
                    analysis_machine_time[component][eq] = group_ready_time + self.analysis_time[workpiece][component]

                # 更新组完成时间
                group_completion_time = group_ready_time + max(
                    self.analysis_time[workpiece][component] for component in components_in_group
                )
                start_analysis_time = group_completion_time

        analysis_completion_time = max(max(machine_times) for machine_times in analysis_machine_time)
        total_time = max(inflation_completion_time, analysis_completion_time)
        return 1.0 / total_time

    def select_parents(self, fitness_values: List[float]) -> Tuple[Dict, Dict]:
        """轮盘赌选择父母"""
        total_fitness = sum(fitness_values)
        probabilities = [f / total_fitness for f in fitness_values]

        parent1_idx = np.random.choice(self.population_size, p=probabilities)
        parent2_idx = np.random.choice(self.population_size, p=probabilities)
        while parent1_idx == parent2_idx:
            parent2_idx = np.random.choice(self.population_size, p=probabilities)

        return self.population[parent1_idx], self.population[parent2_idx]

    def crossover(self, parent1: Dict, parent2: Dict) -> Tuple[Dict, Dict]:
        """交叉操作，确保分析设备列表与组分数一致"""
        if random.random() > self.crossover_rate:
            return parent1.copy(), parent2.copy()

        # 顺序交叉（部分映射交叉）
        def order_crossover(order1, order2):
            size = len(order1)
            a, b = random.sample(range(size), 2)
            if a > b:
                a, b = b, a

            child1 = [-1] * size
            child2 = [-1] * size
            child1[a:b + 1] = order1[a:b + 1]
            child2[a:b + 1] = order2[a:b + 1]

            def fill_child(child, parent, start, end):
                ptr = (end + 1) % size
                for i in range(end + 1, end + 1 + size):
                    current = parent[i % size]
                    if current not in child:
                        child[ptr] = current
                        ptr = (ptr + 1) % size
                return child

            child1 = fill_child(child1, order2, a, b)
            child2 = fill_child(child2, order1, a, b)
            return child1, child2

        # 交叉充气顺序和分析顺序
        child1_inflation_order, child2_inflation_order = order_crossover(
            parent1['inflation_order'], parent2['inflation_order']
        )
        child1_analysis_order, child2_analysis_order = order_crossover(
            parent1['analysis_order'], parent2['analysis_order']
        )

        # 深拷贝确保子染色体独立
        child1 = {
            'inflation_order': child1_inflation_order,
            'inflation_eq': parent1['inflation_eq'].copy(),
            'analysis_order': child1_analysis_order,
            'component_groups': [g.copy() for g in parent1['component_groups']],
            'group_test_order': [g.copy() for g in parent1['group_test_order']],
            'analysis_eq': [g.copy() for g in parent1['analysis_eq']]
        }
        for i in range(len(child1['analysis_eq'])):
            child1['analysis_eq'][i] = [e.copy() for e in child1['analysis_eq'][i]]

        child2 = {
            'inflation_order': child2_inflation_order,
            'inflation_eq': parent2['inflation_eq'].copy(),
            'analysis_order': child2_analysis_order,
            'component_groups': [g.copy() for g in parent2['component_groups']],
            'group_test_order': [g.copy() for g in parent2['group_test_order']],
            'analysis_eq': [g.copy() for g in parent2['analysis_eq']]
        }
        for i in range(len(child2['analysis_eq'])):
            child2['analysis_eq'][i] = [e.copy() for e in child2['analysis_eq'][i]]

        # 充气设备选择交叉
        for i in range(self.num_workpieces):
            if random.random() < 0.5:
                child1['inflation_eq'][i] = parent2['inflation_eq'][i]
                child2['inflation_eq'][i] = parent1['inflation_eq'][i]

        # 组分分组交叉（关键修复：同步校验设备列表长度）
        for i in range(self.num_workpieces):
            if random.random() < 0.5:
                # 交叉组分分组
                child1['component_groups'][i] = parent2['component_groups'][i].copy()
                child2['component_groups'][i] = parent1['component_groups'][i].copy()

                # 重新计算组测试顺序和设备列表
                # 对child1处理
                groups1 = child1['component_groups'][i]
                unique_groups1 = list(sorted(set(groups1)))
                random.shuffle(unique_groups1)
                child1['group_test_order'][i] = unique_groups1
                # 重新生成设备列表确保长度匹配
                new_eq1 = []
                for group_id in unique_groups1:
                    components = [j for j, g in enumerate(groups1) if g == group_id]
                    new_eq1.append([random.randint(0, self.num_analysis_eq_per_component - 1)
                                    for _ in components])
                child1['analysis_eq'][i] = new_eq1

                # 对child2处理
                groups2 = child2['component_groups'][i]
                unique_groups2 = list(sorted(set(groups2)))
                random.shuffle(unique_groups2)
                child2['group_test_order'][i] = unique_groups2
                # 重新生成设备列表确保长度匹配
                new_eq2 = []
                for group_id in unique_groups2:
                    components = [j for j, g in enumerate(groups2) if g == group_id]
                    new_eq2.append([random.randint(0, self.num_analysis_eq_per_component - 1)
                                    for _ in components])
                child2['analysis_eq'][i] = new_eq2

        return child1, child2

    def mutate(self, chromosome: Dict) -> Dict:
        """变异操作，确保组内设备列表与组分数一致"""
        # 深拷贝染色体
        mutated = {
            'inflation_order': chromosome['inflation_order'].copy(),
            'inflation_eq': chromosome['inflation_eq'].copy(),
            'analysis_order': chromosome['analysis_order'].copy(),
            'component_groups': [g.copy() for g in chromosome['component_groups']],
            'group_test_order': [g.copy() for g in chromosome['group_test_order']],
            'analysis_eq': [g.copy() for g in chromosome['analysis_eq']]
        }
        for i in range(len(mutated['analysis_eq'])):
            mutated['analysis_eq'][i] = [e.copy() for e in mutated['analysis_eq'][i]]

        # 充气顺序变异
        if random.random() < self.mutation_rate:
            idx1, idx2 = random.sample(range(self.num_workpieces), 2)
            mutated['inflation_order'][idx1], mutated['inflation_order'][idx2] = \
                mutated['inflation_order'][idx2], mutated['inflation_order'][idx1]

        # 分析顺序变异
        if random.random() < self.mutation_rate:
            idx1, idx2 = random.sample(range(self.num_workpieces), 2)
            mutated['analysis_order'][idx1], mutated['analysis_order'][idx2] = \
                mutated['analysis_order'][idx2], mutated['analysis_order'][idx1]

        # 充气设备选择变异
        for i in range(self.num_workpieces):
            if random.random() < self.mutation_rate:
                mutated['inflation_eq'][i] = random.randint(0, self.num_inflation_eq - 1)

        # 组分分组变异（核心修复：严格同步设备列表）
        for i in range(self.num_workpieces):
            if random.random() < self.mutation_rate:
                # 随机改变一个组分的组号
                comp_idx = random.randint(0, self.num_components - 1)
                current_groups = mutated['component_groups'][i]
                possible_groups = list(set(current_groups))

                if len(possible_groups) == 1:
                    new_group = 1 if possible_groups[0] == 0 else 0
                else:
                    new_group = random.choice([g for g in possible_groups if g != current_groups[comp_idx]])

                current_groups[comp_idx] = new_group

                # 更新组测试顺序和设备列表
                unique_groups = list(sorted(set(current_groups)))
                random.shuffle(unique_groups)
                mutated['group_test_order'][i] = unique_groups

                # 重新生成分析设备选择（确保与组内组分数一致）
                new_analysis_eq = []
                for group_id in unique_groups:
                    components_in_group = [j for j, g in enumerate(current_groups) if g == group_id]
                    # 设备列表长度严格等于组内组分数
                    new_analysis_eq.append([
                        random.randint(0, self.num_analysis_eq_per_component - 1)
                        for _ in components_in_group
                    ])
                mutated['analysis_eq'][i] = new_analysis_eq

        # 分析设备选择变异（只在有效索引范围内操作）
        for i in range(self.num_workpieces):
            for group_idx in range(len(mutated['group_test_order'][i])):
                if random.random() < self.mutation_rate:
                    eq_list = mutated['analysis_eq'][i][group_idx]
                    if len(eq_list) > 0:
                        comp_idx = random.randint(0, len(eq_list) - 1)
                        eq_list[comp_idx] = random.randint(0, self.num_analysis_eq_per_component - 1)

        return mutated

    def evolve(self) -> Tuple[Dict, float]:
        """执行遗传算法进化过程"""
        self.initialize_population()
        best_fitness = 0
        best_chromosome = None

        for generation in range(self.generations):
            try:
                fitness_values = [self.calculate_fitness(chrom) for chrom in self.population]
            except Exception as e:
                print(f"计算适应度时出错: {e}")
                # 出错时重新生成种群
                self.initialize_population()
                fitness_values = [self.calculate_fitness(chrom) for chrom in self.population]

            current_best_idx = np.argmax(fitness_values)
            current_best_fitness = fitness_values[current_best_idx]
            current_best_chromosome = self.population[current_best_idx]

            if current_best_fitness > best_fitness:
                best_fitness = current_best_fitness
                best_chromosome = current_best_chromosome

            if generation % 10 == 0:
                print(
                    f"Generation {generation}: Best Fitness = {current_best_fitness:.4f}, Total Time = {1 / current_best_fitness:.2f}")

            # 生成新一代种群
            new_population = []
            elite_size = max(1, int(self.population_size * 0.1))
            elite_indices = np.argsort(fitness_values)[-elite_size:]
            for idx in elite_indices:
                new_population.append(self.population[idx])

            while len(new_population) < self.population_size:
                parent1, parent2 = self.select_parents(fitness_values)
                child1, child2 = self.crossover(parent1, parent2)
                child1 = self.mutate(child1)
                child2 = self.mutate(child2)
                new_population.append(child1)
                if len(new_population) < self.population_size:
                    new_population.append(child2)

            self.population = new_population

        return best_chromosome, best_fitness


# 示例用法
if __name__ == "__main__":
    num_workpieces = 5  # 5个气瓶
    num_components = 3  # 每个气瓶有3种组分
    num_inflation_eq = 2  # 2台充气设备
    num_analysis_eq = 2  # 每种组分有2台分析设备

    ga = GasSchedulingGA(
        num_workpieces=num_workpieces,
        num_components=num_components,
        num_inflation_eq=num_inflation_eq,
        num_analysis_eq_per_component=num_analysis_eq,
        population_size=50,
        generations=100,
        crossover_rate=0.8,
        mutation_rate=0.1
    )

    best_solution, best_fitness = ga.evolve()

    print("\nBest Solution Found:")
    print(f"Total Time: {1 / best_fitness:.2f}")
    print("Inflation Order:", best_solution['inflation_order'])
    print("Inflation Equipment:", best_solution['inflation_eq'])
    print("Analysis Order:", best_solution['analysis_order'])

    print("\nComponent Groups (each row is a workpiece, numbers indicate group):")
    for i, groups in enumerate(best_solution['component_groups']):
        print(f"Workpiece {i}: {groups}")
